import socket
import subprocess
import time
import traceback

import socks

from absl import flags, app

from coin.exchange.bithumb.kr_rest_web.order_client import BithumbWebOrder

flags.DEFINE_string("cookie_file", None, "Path of cookie_file which has authentication token.")

flags.DEFINE_string("remote_host_string", None, "SSH host string of proxy server.")

flags.DEFINE_string("coin_deploy_dir", "./../../coin_deploy/", "")

flags.DEFINE_string("ssh_config_file", "./kr_ssh_config/ssh_config", "")

flags.DEFINE_integer("retry_count", 5, "")

FLAGS = flags.FLAGS


def _find_free_port():
  s = socket.socket()
  s.bind(('', 0))  # Bind to a free port provided by the host.
  port = s.getsockname()[1]  # Return the port number assigned.
  return port


def dynamic_port_forwarding(remote_host_string,
                            ssh_config_file=None,
                            bind_address='127.0.0.1',
                            port=None,
                            coin_deploy_dir=None):
  """Return (Subprocess, port)_pair when it successfully creates a tunnel
  otherwise returns None."""
  cmds = ['ssh', '-o', 'StrictHostKeyChecking=no', '-o', 'UserKnownHostsFile=/dev/null']

  if ssh_config_file:
    cmds.extend(['-F', ssh_config_file])

  if port is None:
    port = _find_free_port()

  cmds.append(remote_host_string)
  cmds.extend(['-ND', '%s:%d' % (bind_address, port)])

  print("port forwarding cmd:", cmds)
  kwargs = dict()
  if coin_deploy_dir:
    kwargs['cwd'] = coin_deploy_dir
  popen = subprocess.Popen(cmds, **kwargs)

  return (popen, port)


def main(_):
  assert FLAGS.cookie_file

  for i in range(FLAGS.retry_count):
    popen = None
    try:
      popen, port = dynamic_port_forwarding(FLAGS.remote_host_string,
                                            FLAGS.ssh_config_file,
                                            coin_deploy_dir=FLAGS.coin_deploy_dir)
      time.sleep(10)
      if popen.poll() is None:
        pass

      proxy_server = 'socks5://127.0.0.1:%d' % port
      print("Proxy launched on : ", proxy_server)

      socks.setdefaultproxy(socks.PROXY_TYPE_SOCKS5, 'localhost', port)
      socket.socket = socks.socksocket

      # Querying balance raises an exception when cookie is not valid.
      print(BithumbWebOrder('XRP', FLAGS.cookie_file).query_balance())
    except Exception:
      if i + 1 == FLAGS.retry_count:
        raise
      else:
        traceback.print_exc()
        time.sleep(5)
    else:
      break
    finally:
      if popen:
        popen.kill()

  return 0


if __name__ == '__main__':
  app.run(main)
